# https://docs.fenicsproject.org/dolfinx/v0.9.0/python/demos/demo_navier-stokes.html
# 源代码
import importlib.util

if importlib.util.find_spec("petsc4py") is not None:
    import dolfinx

    if not dolfinx.has_petsc:
        print("This demo requires DOLFINx to be compiled with PETSc enabled.")
        exit(0)
else:
    print("This demo requires petsc4py.")
    exit(0)

from mpi4py import MPI

# +
import numpy as np

from dolfinx import default_real_type, fem, io, mesh
from dolfinx.fem.petsc import assemble_matrix_block, assemble_vector_block
from ufl import (
    CellDiameter,
    FacetNormal,
    MixedFunctionSpace,
    TestFunctions,
    TrialFunctions,
    avg,
    conditional,
    div,
    dot,
    dS,  # dS 相邻单元的共享面 【内部界面 非物理边界】
    ds,  # 计算区域的暴露面  【外部边界 物理边界】
    dx,
    extract_blocks,
    grad,
    gt,
    inner,
    outer,
)

try:
    from petsc4py import PETSc

    import dolfinx

    if not dolfinx.has_petsc:
        print("This demo requires DOLFINx to be compiled with PETSc enabled.")
        exit(0)
except ModuleNotFoundError:
    print("This demo requires petsc4py.")
    exit(0)


if np.issubdtype(PETSc.ScalarType, np.complexfloating):  # type: ignore
    print("Demo should only be executed with DOLFINx real mode")
    exit(0)
# -


# We also define some helper functions that will be used later


# +
def norm_L2(comm, v):
    """Compute the L2(Ω)-norm of v"""
    return np.sqrt(comm.allreduce(fem.assemble_scalar(fem.form(inner(v, v) * dx)), op=MPI.SUM))



# 在并行计算环境下，计算函数 v 在网格 msh 所定义的整个物理区域上的平均值，核心逻辑是 “函数积分除以区域体积”。
def domain_average(msh, v):
    """Compute the average of a function over the domain"""
    vol = msh.comm.allreduce(
        fem.assemble_scalar(fem.form(fem.Constant(msh, default_real_type(1.0)) * dx)), op=MPI.SUM
    )
    return (1 / vol) * msh.comm.allreduce(fem.assemble_scalar(fem.form(v * dx)), op=MPI.SUM)


def u_e_expr(x):
    """Expression for the exact velocity solution to Kovasznay flow"""
    return np.vstack(
        (
            1
            - np.exp((Re / 2 - np.sqrt(Re**2 / 4 + 4 * np.pi**2)) * x[0])
            * np.cos(2 * np.pi * x[1]),
            (Re / 2 - np.sqrt(Re**2 / 4 + 4 * np.pi**2))
            / (2 * np.pi)
            * np.exp((Re / 2 - np.sqrt(Re**2 / 4 + 4 * np.pi**2)) * x[0])
            * np.sin(2 * np.pi * x[1]),
        )
    )


def p_e_expr(x):
    """Expression for the exact pressure solution to Kovasznay flow"""
    return (1 / 2) * (1 - np.exp(2 * (Re / 2 - np.sqrt(Re**2 / 4 + 4 * np.pi**2)) * x[0]))


def f_expr(x):
    """Expression for the applied force"""
    return np.vstack((np.zeros_like(x[0]), np.zeros_like(x[0])))


# -

# We define some simulation parameters


n = 16
num_time_steps = 25
t_end = 10
Re = 25  # Reynolds Number
k = 1  # Polynomial degree

# Next, we create a mesh and the required functions spaces over it.
# Since the velocity uses an $H(\text{div})$-conforming function
# space, we also create a vector valued discontinuous Lagrange space to
# interpolate into for artifact free visualisation.

# +
msh = mesh.create_unit_square(MPI.COMM_WORLD, n, n)

# Function spaces for the velocity and for the pressure
V = fem.functionspace(msh, ("Raviart-Thomas", k + 1))
Q = fem.functionspace(msh, ("Discontinuous Lagrange", k))
VQ = MixedFunctionSpace(V, Q)

# Funcion space for visualising the velocity field
gdim = msh.geometry.dim
W = fem.functionspace(msh, ("Discontinuous Lagrange", k + 1, (gdim,)))

# Define trial and test functions

u, p = TrialFunctions(VQ)
v, q = TestFunctions(VQ)

delta_t = fem.Constant(msh, default_real_type(t_end / num_time_steps))
alpha = fem.Constant(msh, default_real_type(6.0 * k**2))

h = CellDiameter(msh)
n = FacetNormal(msh)

# phi：分片光滑的向量函数（如速度场、应力张量等），在 DG 方法中，它在单元内部连续，但在单元边界可能间断。
# n：单元边界的外法向量场。
def jump(phi, n):
    # outer: 外积
    return outer(phi("+"), n("+")) + outer(phi("-"), n("-"))


# -

# 求解initial condition 忽略非线性项（对流项）


# +
# a的前四项对应a(u, v)
a = (1.0 / Re) * (
    inner(grad(u), grad(v)) * dx
    - inner(avg(grad(u)), jump(v, n)) * dS
    - inner(jump(u, n), avg(grad(v))) * dS
    # 这里是a_h(u,v)的最后一项 h_K表示单元K的特征尺寸 通常去单元
    # h = CellDiameter(msh) 计算网格中每个单元特征尺寸的内置函数
    # 这里和旁边的单元直径做了一下平均
    + (alpha / avg(h)) * inner(jump(u, n), jump(v, n)) * dS
    
    - inner(grad(u), outer(v, n)) * ds
    - inner(outer(u, n), grad(v)) * ds
    # L_ah(v_h)的第二项
    + (alpha / h) * inner(outer(u, n), outer(v, n)) * ds
)
# 非线性项c(w;u,v)放到后面while循环里

# b_h(v_h, p_h)
a -= inner(p, div(v)) * dx
# b_h(u_h, q_h)
a -= inner(div(u), q) * dx

# extract_blocks用于把定义在混合函数空间的形式a 提取出对应空间的分块
a_blocked = fem.form(extract_blocks(a))

f = fem.Function(W)

u_D = fem.Function(V)
u_D.interpolate(u_e_expr)
# (f, v) + 
L = inner(f, v) * dx + (1 / Re) * (
    # L_ah(v_h)
    -inner(outer(u_D, n), grad(v)) * ds + (alpha / h) * inner(outer(u_D, n), outer(v, n)) * ds
)
L += inner(fem.Constant(msh, default_real_type(0.0)), q) * dx
L_blocked = fem.form(extract_blocks(L))

# Boundary conditions
# exterior_facet_indices: 计算当前进程所 “拥有” 的 外部面（exterior facets）的索引
# 之前的写法是定义网格边界的判定函数 利用几何关系
# 现在利用网格的拓扑信息 比较擅长处理复杂的形状 但无法精确指定边界的一部分（比如矩形只想设置左右边界 就用之前的写法更方便）
boundary_facets = mesh.exterior_facet_indices(msh.topology)
boundary_vel_dofs = fem.locate_dofs_topological(V, msh.topology.dim - 1, boundary_facets)
bc_u = fem.dirichletbc(u_D, boundary_vel_dofs)
bcs = [bc_u]

# Assemble Stokes problem
A = assemble_matrix_block(a_blocked, bcs=bcs)
A.assemble()
b = assemble_vector_block(L_blocked, a_blocked, bcs=bcs)

# Create and configure solver
ksp = PETSc.KSP().create(msh.comm)  # type: ignore
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
opts = PETSc.Options()  # type: ignore
opts["mat_mumps_icntl_14"] = 80  # Increase MUMPS working memory
opts["mat_mumps_icntl_24"] = 1  # Option to support solving a singular matrix (pressure nullspace)
opts["mat_mumps_icntl_25"] = 0  # Option to support solving a singular matrix (pressure nullspace)
opts["ksp_error_if_not_converged"] = 1
ksp.setFromOptions()

# Solve Stokes for initial condition
x = A.createVecRight()
try:
    ksp.solve(b, x)
except PETSc.Error as e:  # type: ignore
    if e.ierr == 92:
        print("The required PETSc solver/preconditioner is not available. Exiting.")
        print(e)
        exit(0)
    else:
        raise e

# Split the solution
u_h = fem.Function(V)
p_h = fem.Function(Q)
p_h.name = "p"
# 速度的总本地自由度
offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
u_h.x.array[:offset] = x.array_r[:offset]
u_h.x.scatter_forward()
p_h.x.array[: (len(x.array_r) - offset)] = x.array_r[offset:]
p_h.x.scatter_forward()
# Subtract the average of the pressure since it is only determined up to
# a constant 减去了整个计算区域的压力平均值
p_h.x.array[:] -= domain_average(msh, p_h)

u_vis = fem.Function(W)
u_vis.name = "u"
u_vis.interpolate(u_h)

# Write initial condition to file
t = 0.0
try:
    u_file = io.VTXWriter(msh.comm, "u.bp", u_vis)
    p_file = io.VTXWriter(msh.comm, "p.bp", p_h)
    u_file.write(t)
    p_file.write(t)
except AttributeError:
    print("File output requires ADIOS2.")

# Create function to store solution and previous time step
u_n = fem.Function(V)
u_n.x.array[:] = u_h.x.array
# -

# Now we add the time stepping and convective terms

# +
# conditional函数是 条件表达式构造器
# gt: greater than 左侧大于右侧的符号表达式
lmbda = conditional(gt(dot(u_n, n), 0), 1, 0)
# u_uw是迎风速度 从信息上游取值
# +表示当前单元的界面侧 -表示相邻单元的界面侧

u_uw = lmbda("+") * u("+") + lmbda("-") * u("-")
a += (
    inner(u / delta_t, v) * dx
    - inner(u, div(outer(v, u_n))) * dx
    + inner((dot(u_n, n))("+") * u_uw, v("+")) * dS
    + inner((dot(u_n, n))("-") * u_uw, v("-")) * dS
    + inner(dot(u_n, n) * lmbda * u, v) * ds
)
a_blocked = fem.form(extract_blocks(a))

L += inner(u_n / delta_t, v) * dx - inner(dot(u_n, n) * (1 - lmbda) * u_D, v) * ds
L_blocked = fem.form(extract_blocks(L))

# Time stepping loop
for n in range(num_time_steps):
    t += delta_t.value

    A.zeroEntries()  # 清空矩阵
    fem.petsc.assemble_matrix_block(A, a_blocked, bcs=bcs)  # type: ignore
    A.assemble()

    with b.localForm() as b_loc:  # 获取向量的本地部分
        b_loc.set(0)
    fem.petsc.assemble_vector_block(b, L_blocked, a_blocked, bcs=bcs)  # type: ignore

    # Compute solution
    ksp.solve(b, x)

    u_h.x.array[:offset] = x.array_r[:offset]
    u_h.x.scatter_forward()
    p_h.x.array[: (len(x.array_r) - offset)] = x.array_r[offset:]
    p_h.x.scatter_forward()
    p_h.x.array[:] -= domain_average(msh, p_h)

    u_vis.interpolate(u_h)

    # Write to file
    try:
        u_file.write(t)
        p_file.write(t)
    except NameError:
        pass

    # Update u_n
    u_n.x.array[:] = u_h.x.array

try:
    u_file.close()
    p_file.close()
except NameError:
    pass
# -

# Now we compare the computed solution to the exact solution

# +
# Function spaces for exact velocity and pressure
# 把精确解插值到高次空间 使得插值误差远小于数值解的误差
# 精确解通常光滑 所以插值到lagrange空间
V_e = fem.functionspace(msh, ("Lagrange", k + 3, (gdim,)))
Q_e = fem.functionspace(msh, ("Lagrange", k + 2))

u_e = fem.Function(V_e)
u_e.interpolate(u_e_expr)

p_e = fem.Function(Q_e)
p_e.interpolate(p_e_expr)

# Compute errors
e_u = norm_L2(msh.comm, u_h - u_e)
e_div_u = norm_L2(msh.comm, div(u_h))

# This scheme conserves mass exactly, so check this
assert np.isclose(e_div_u, 0.0, atol=float(1.0e5 * np.finfo(default_real_type).eps))
p_e_avg = domain_average(msh, p_e)
e_p = norm_L2(msh.comm, p_h - (p_e - p_e_avg))

if msh.comm.rank == 0:
    print(f"e_u = {e_u}")
    print(f"e_div_u = {e_div_u}")
    print(f"e_p = {e_p}")
# -
